/**
 * 
 */
package edu.umd.clip.lm.model.training.metrics;

import org.w3c.dom.Element;

import edu.umd.clip.lm.model.training.ContextVariable;
import edu.umd.clip.lm.model.training.ContextVariableStats;
import edu.umd.clip.lm.util.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class DistanceMetrics extends ContextVariableMetrics {

	public DistanceMetrics() {
		super("Distance-based metrics: (H(w|x) + H(x|w)) / H(w,x)");
	}
	
	/**
	 * @param elem
	 */
	public DistanceMetrics(Element elem) {
		super(elem);
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.metrics.ContextVariableMetrics#computeScore(edu.umd.clip.lm.model.training.ContextVariable, edu.umd.clip.lm.model.training.ContextVariableStats)
	 */
	@Override
	public double computeScore(ContextVariable ctxVar, ContextVariableStats stat) {
		// (H(w|x) + H(x|w)) / H(w,x) = 1 - I(x;w) / H(w,x) = 1 - (H(w) - H(w|x))/H(w,x)
		
		long totalCount = 0;
		double condEntropy = 0;
		double jointEntropy = 0;
		for(Object d : stat.x2wDistributions.values()) {
			ConstCountDistribution dist = (ConstCountDistribution) d;
			
			long keys[] = dist.keys();
			long counts[] = dist.values();
			for(int i=0; i<keys.length; ++i) {
				double element = - counts[i] * ProbMath.log2(counts[i]);
				condEntropy += element;
				jointEntropy += element;
			}

			condEntropy += dist.getTotalCount() * ProbMath.log2(dist.getTotalCount());
			totalCount += dist.getTotalCount();
		}
		condEntropy /= totalCount;
		
		jointEntropy /= totalCount;
		jointEntropy += ProbMath.log2(totalCount);
		
		double wordEntropy = 0;
		for(long count : stat.wordCounts.values()) {
			wordEntropy -= count * ProbMath.log2(count);
		}
		wordEntropy /= totalCount;
		wordEntropy += ProbMath.log2(totalCount);
		
		double distance = 1 - (wordEntropy - condEntropy) / jointEntropy;
		return distance;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.metrics.ContextVarStatUser#needContextVarCounts()
	 */
	@Override
	public boolean needContextVarCounts() {
		return false;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.metrics.ContextVarStatUser#needVarToWordCounts()
	 */
	@Override
	public boolean needVarToWordCounts() {
		return true;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.metrics.ContextVarStatUser#needWordCounts()
	 */
	@Override
	public boolean needWordCounts() {
		return true;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.metrics.ContextVarStatUser#needWordToVarCounts()
	 */
	@Override
	public boolean needWordToVarCounts() {
		return true;
	}

}
